#批量处理，脚本，阻塞列队，一个一个跑以免机器吃不消
#首先在input文件夹里面，创建批处理文件夹，把这批量跑的fasta文件全放进去
#输入 python one_by_one.py <文件夹名>即可
#作者：zhuxf


import sys
import shutil
import os
import numpy
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker


filePath = sys.argv[1]
#较小的遗传数据库配置 (reduced_dbs) 或完整的遗传数据库配置 (full_dbs)（默认值：' full_dbs '）
#db_preset = sys.argv[2]
db_preset = "reduced_dbs"

def run_af2(filePath,db_preset):
    filelist_temp = os.listdir("./input/" + filePath)
    filelist = []

    # 把奇怪后缀全部换成.fasta
    for item in filelist_temp:
        filelist.append(item.split(".")[0] + ".fasta")

    # 按照数字排序
    filelist.sort()
    # print(filelist)
    filelist3 = []
    filelist4 = []
    # models = [1, 2, 3, 4, 5]
    # print("预备执行========================alphafold=====================创建文件夹：single-{output_dir}".format(output_dir=filePath))
    # 创建批处理的输出文件夹
    os.system("mkdir ./output/single-{output_dir}".format(output_dir=filePath))

    # 删除无用项目文件夹
    filelist_detect = os.listdir("./output/single-{output_dir}/".format(output_dir=filePath))

    for i in filelist_detect:
        if not os.path.exists("./output/single-{output_dir}/{output_file}/{output_file}".format(output_dir=filePath,output_file=i)):
            shutil.rmtree("./output/single-{output_dir}/{output_file}".format(output_dir=filePath, output_file=i))

    filelist2 = os.listdir("./output/single-{output_dir}/".format(output_dir=filePath))

    for i in filelist2:
        filelist3.append(i + ".fasta")
    # 差集
    for i in filelist:
        if i not in filelist3:
            filelist4.append(i)

    print("开始执行任务列队：======>")
    print(filelist4)

    item = filelist4[-1]
    # print("bash run_alphafold.sh -d ../alphafold_database -o ./output/批处理-{output_dir} -m model_1 -f ./input/{input} -t 2000-01-10".format(output_dir=filePath,input=item))
    print("开始执行===================alphafold===========================文件名：" + item)
    # 开始跑alphafold
    os.system(
        "bash run_alphafold.sh -d ../alphafold_database -o ./output/single-{output_dir} -f ./input/{dir}/{input} -t 2020-05-14 -c {db_preset} -r false".format(output_dir=filePath, dir=filePath, input=item, db_preset=db_preset))
    # 修改权限
    os.system("sudo chmod 777 ./output/single-{output_dir}/{output_file}".format(output_dir=filePath,output_file=item.split(".")[0]))
    os.system("sudo mkdir ./output/single-{output_dir}/{output_file}/{output_file}".format(output_dir=filePath,output_file=item.split(".")[0]))
    os.system("sudo chmod 777 ./output/single-{output_dir}/{output_file}/{output_file}".format(output_dir=filePath,output_file=item.split(".")[0]))
    os.system("sudo cp ./output/single-{output_dir}/{output_file}/rank* ./output/single-{output_dir}/{output_file}/{output_file}".format(output_dir=filePath, output_file=item.split(".")[0]))
    #画图
    draw_path = "./output/single-{output_dir}/{output_file}/{output_file}".format(output_dir=filePath,output_file=item.split(".")[0])
    draw_single(draw_path)
    draw_merge(draw_path)

    if len(filelist4[0:-1])>0:
        run_af2(filePath,db_preset)


def draw_single(draw_path):
    draw_path = draw_path
    def savepic(names,draw_path):
        file_tmp = open("{draw_path}/{names}".format(names=names,draw_path=draw_path), 'r')
        tmp_lines = [i for i in file_tmp]
        file_tmp.close()
        aa_number = []
        plddt = []
        tmp_lines_clean = []

        for i in tmp_lines:
            if i[0:4] == "ATOM":
                tmp_lines_clean.append(i)

        file_tmp = open("{draw_path}/{names}_clean.pdb".format(names=names,draw_path=draw_path), 'w+')
        file_tmp.writelines(tmp_lines_clean)
        file_tmp.close()

        file_tmp = open("{draw_path}/{names}_clean.pdb".format(names=names,draw_path=draw_path), 'r')
        tmp_lines = [i for i in file_tmp]
        file_tmp.close()

        for i in range(len(tmp_lines)):
            if int(tmp_lines[i][22:30]) not in aa_number:
                aa_number.append(int(tmp_lines[i][22:30]))
                plddt.append(float(tmp_lines[i][61:66]))
        dict = {'aa_number': aa_number, 'plddt': plddt}

        df = pd.DataFrame(dict)
        print(df)

        x = df['aa_number']
        y1 = df['plddt']

        # Plot Line1 (Left Y Axis)
        fig, ax1 = plt.subplots(1, 1, figsize=(16, 9), dpi=600)
        ax1.plot(x, y1, color='tab:red')

        # Decorations
        # ax1 (left Y axis)
        ax1.set_xlabel('Amino acid number', fontsize=24)
        # ax1.tick_params(axis='x', rotation=0, labelsize=10,labelcolor='tab:red')
        ax1.tick_params(axis='x', rotation=0, labelsize=10)
        # 显示所有横坐标
        # plt.xticks(x)
        ax1.xaxis.set_major_locator(ticker.MultipleLocator(base=15))

        ax1.set_ylabel('Plddt(0~100)', fontsize=24)
        # ax1.tick_params(axis='y', rotation=0, labelcolor='tab:red' )
        ax1.tick_params(axis='y', rotation=0)
        ax1.grid(alpha=.2)

        fig.tight_layout()
        plt.savefig("{draw_path}/{names}.png".format(names=names,draw_path=draw_path))

    list = ["ranked_0.pdb", "ranked_1.pdb", "ranked_2.pdb", "ranked_3.pdb", "ranked_4.pdb"]
    for n in list:
        savepic(n,draw_path)

    for i in range(5):
        if os.path.exists("{draw_path}/ranked_{i}.pdb_clean.pdb".format(i=i,draw_path=draw_path)):
            os.remove("{draw_path}/ranked_{i}.pdb_clean.pdb".format(i=i,draw_path=draw_path))

def draw_merge(draw_path):

    file_tmp1 = open("{draw_path}/ranked_0.pdb".format(draw_path=draw_path), 'r')
    file_tmp2 = open("{draw_path}/ranked_1.pdb".format(draw_path=draw_path), 'r')
    file_tmp3 = open("{draw_path}/ranked_2.pdb".format(draw_path=draw_path), 'r')
    file_tmp4 = open("{draw_path}/ranked_3.pdb".format(draw_path=draw_path), 'r')
    file_tmp5 = open("{draw_path}/ranked_4.pdb".format(draw_path=draw_path), 'r')
    tmp_lines1 = [i for i in file_tmp1]
    tmp_lines2 = [i for i in file_tmp2]
    tmp_lines3 = [i for i in file_tmp3]
    tmp_lines4 = [i for i in file_tmp4]
    tmp_lines5 = [i for i in file_tmp5]
    file_tmp1.close()
    file_tmp2.close()
    file_tmp3.close()
    file_tmp4.close()
    file_tmp5.close()
    tmp_lines1_clean = []
    tmp_lines2_clean = []
    tmp_lines3_clean = []
    tmp_lines4_clean = []
    tmp_lines5_clean = []

    for i in tmp_lines1:
        if i[0:4] == "ATOM":
            tmp_lines1_clean.append(i)

    for i in tmp_lines2:
        if i[0:4] == "ATOM":
            tmp_lines2_clean.append(i)

    for i in tmp_lines3:
        if i[0:4] == "ATOM":
            tmp_lines3_clean.append(i)

    for i in tmp_lines4:
        if i[0:4] == "ATOM":
            tmp_lines4_clean.append(i)

    for i in tmp_lines5:
        if i[0:4] == "ATOM":
            tmp_lines5_clean.append(i)

    file_tmp1 = open("{}/ranked_0_clean.pdb".format(draw_path), 'w+')
    file_tmp1.writelines(tmp_lines1_clean)
    file_tmp1.close()
    file_tmp2 = open("{}/ranked_1_clean.pdb".format(draw_path), 'w+')
    file_tmp2.writelines(tmp_lines2_clean)
    file_tmp2.close()
    file_tmp3 = open("{}/ranked_2_clean.pdb".format(draw_path), 'w+')
    file_tmp3.writelines(tmp_lines3_clean)
    file_tmp3.close()
    file_tmp4 = open("{}/ranked_3_clean.pdb".format(draw_path), 'w+')
    file_tmp4.writelines(tmp_lines4_clean)
    file_tmp4.close()
    file_tmp5 = open("{}/ranked_4_clean.pdb".format(draw_path), 'w+')
    file_tmp5.writelines(tmp_lines5_clean)
    file_tmp5.close()

    file_tmp1 = open("{}/ranked_0_clean.pdb".format(draw_path), 'r')
    file_tmp2 = open("{}/ranked_1_clean.pdb".format(draw_path), 'r')
    file_tmp3 = open("{}/ranked_2_clean.pdb".format(draw_path), 'r')
    file_tmp4 = open("{}/ranked_3_clean.pdb".format(draw_path), 'r')
    file_tmp5 = open("{}/ranked_4_clean.pdb".format(draw_path), 'r')
    tmp_lines1 = [i for i in file_tmp1]
    tmp_lines2 = [i for i in file_tmp2]
    tmp_lines3 = [i for i in file_tmp3]
    tmp_lines4 = [i for i in file_tmp4]
    tmp_lines5 = [i for i in file_tmp5]

    file_tmp1.close()
    file_tmp2.close()
    file_tmp3.close()
    file_tmp4.close()
    file_tmp5.close()

    aa_number = []
    plddt1 = []
    plddt2 = []
    plddt3 = []
    plddt4 = []
    plddt5 = []

    for i in range(len(tmp_lines1)):
        if int(tmp_lines1[i][22:30]) not in aa_number:
            aa_number.append(int(tmp_lines1[i][22:30]))
            plddt1.append(float(tmp_lines1[i][61:66]))
    aa_number = []
    for i in range(len(tmp_lines2)):
        if int(tmp_lines2[i][22:30]) not in aa_number:
            aa_number.append(int(tmp_lines2[i][22:30]))
            plddt2.append(float(tmp_lines2[i][61:66]))
    aa_number = []
    for i in range(len(tmp_lines3)):
        if int(tmp_lines3[i][22:30]) not in aa_number:
            aa_number.append(int(tmp_lines3[i][22:30]))
            plddt3.append(float(tmp_lines3[i][61:66]))
    aa_number = []
    for i in range(len(tmp_lines4)):
        if int(tmp_lines4[i][22:30]) not in aa_number:
            aa_number.append(int(tmp_lines4[i][22:30]))
            plddt4.append(float(tmp_lines4[i][61:66]))
    aa_number = []
    for i in range(len(tmp_lines5)):
        if int(tmp_lines5[i][22:30]) not in aa_number:
            aa_number.append(int(tmp_lines5[i][22:30]))
            plddt5.append(float(tmp_lines5[i][61:66]))
    dict1 = {'aa_number': aa_number, 'plddt': plddt1}
    dict2 = {'aa_number': aa_number, 'plddt': plddt2}
    dict3 = {'aa_number': aa_number, 'plddt': plddt3}
    dict4 = {'aa_number': aa_number, 'plddt': plddt4}
    dict5 = {'aa_number': aa_number, 'plddt': plddt5}

    df1 = pd.DataFrame(dict1)
    df2 = pd.DataFrame(dict2)
    df3 = pd.DataFrame(dict3)
    df4 = pd.DataFrame(dict4)
    df5 = pd.DataFrame(dict5)
    print(df1)
    print(df2)
    print(df3)
    print(df4)
    print(df5)

    x = df1['aa_number']
    y1 = df1['plddt']
    y2 = df2['plddt']
    y3 = df3['plddt']
    y4 = df4['plddt']
    y5 = df5['plddt']

    # Plot Line1 (Left Y Axis)
    fig, ax1 = plt.subplots(1, 1, figsize=(16, 9), dpi=600)

    ax1.plot(x, y5, color='tab:orange', label="ranked_4")
    ax1.plot(x, y4, color='tab:purple', label="ranked_3")
    ax1.plot(x, y3, color='tab:blue', label="ranked_2")
    ax1.plot(x, y2, color='tab:green', label="ranked_1")
    ax1.plot(x, y1, color='tab:red', label="ranked_0")
    plt.legend()
    # Decorations
    # ax1 (left Y axis)
    ax1.set_xlabel('Amino acid number', fontsize=24)
    # ax1.tick_params(axis='x', rotation=0, labelsize=10,labelcolor='tab:red')
    ax1.tick_params(axis='x', rotation=0, labelsize=10)
    # 显示所有横坐标
    # plt.xticks(x)
    ax1.xaxis.set_major_locator(ticker.MultipleLocator(base=30))

    ax1.set_ylabel('Plddt(0~100)', fontsize=24)
    # ax1.tick_params(axis='y', rotation=0, labelcolor='tab:red' )
    ax1.tick_params(axis='y', rotation=0)
    ax1.grid(alpha=.1)

    fig.tight_layout()

    plt.savefig('{}/ranked_merge.png'.format(draw_path))
    # plt.show()

    for i in range(5):
        if os.path.exists("{draw_path}/ranked_{i}_clean.pdb".format(i=i,draw_path=draw_path)):
            os.remove("{draw_path}/ranked_{i}_clean.pdb".format(i=i,draw_path=draw_path))



run_af2(filePath,db_preset)